#ifndef _DatasetIO_dataset_io_panovideo_H_
#define _DatasetIO_dataset_io_panovideo_H_

#include "dataset_io.h"
#include "dataset_io_utils.h"

namespace DatasetIO {

class PanoVideoDataset : public Dataset {
    std::string path;

    std::vector<double> image_timestamps;
    std::shared_ptr<cv::VideoCapture> pano_vicap_ptr;
    size_t cur_image_id = 0;

    // Imu
    std::vector<double> imu_timestamps;
    std::vector<AccelData::Ptr> accel_data;
    std::vector<GyroData::Ptr> gyro_data;
    size_t cur_imu_id = 0;

public:
    PanoVideoDataset() : Dataset() { type = DatasetType::kPanoVideo; }
    ~PanoVideoDataset(){};

    std::vector<double> &get_image_timestamps() { return image_timestamps; }

    const std::vector<AccelData::Ptr> &get_accel_data() const {
        return accel_data;
    }
    const std::vector<GyroData::Ptr> &get_gyro_data() const {
        return gyro_data;
    }

    MonoImageData::Ptr get_next_mono_image_data() {
        MonoImageData::Ptr res;
        cv::Mat image;

        if (pano_vicap_ptr->read(image)) {
            double t_s = image_timestamps[cur_image_id++];
            res = std::make_shared<MonoImageData>();
            res->timestamp = t_s;
            res->data = image.clone();
        }

        return res;
    }

    std::vector<AccelData::Ptr> get_range_accel_data(
        const double &start_timestamp, const double &end_timestamp) {
        std::vector<AccelData::Ptr> res;

        auto index = find_closest_index(imu_timestamps, 1, start_timestamp);
        if (index.size() != 0) {
            size_t start_index = index[0];
            while (imu_timestamps[start_index] <= end_timestamp) {
                if (imu_timestamps[start_index] > start_timestamp) {
                    res.push_back(accel_data[start_index]);
                }
                start_index++;
            }

            cur_imu_id = start_index;
        }

        return res;
    }

    std::vector<GyroData::Ptr> get_range_gyro_data(
        const double &start_timestamp, const double &end_timestamp) {
        std::vector<GyroData::Ptr> res;

        auto index = find_closest_index(imu_timestamps, 1, start_timestamp);
        if (index.size() != 0) {
            size_t start_index = index[0];
            while (imu_timestamps[start_index] <= end_timestamp) {
                if (imu_timestamps[start_index] > start_timestamp) {
                    res.push_back(gyro_data[start_index]);
                }
                start_index++;
            }

            cur_imu_id = start_index;
        }

        return res;
    }

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    friend class PanoVideoIO;
};

class PanoVideoIO : public DatasetIoInterface {
public:
    PanoVideoIO() {}

    void read(const std::string &path) {
        if (!fs::exists(path))
            LOG(WARNING) << "No dataset found in " << path;

        data.reset(new PanoVideoDataset);

        data->path = path;
        data->pano_vicap_ptr = std::make_shared<cv::VideoCapture>(
            path + "/video.mp4", cv::CAP_FFMPEG);
        data->cur_image_id = 0;

        read_image_timestamps(path + "/image_times.txt");
        if (fs::exists(path + "/imu_data.txt"))
            read_imu_data(path + "/imu_data.txt");
    }

    void reset() { data.reset(); }

    Dataset::Ptr get_data() { return data; }

private:
    void read_image_timestamps(const std::string &path) {
        std::ifstream f(path);
        std::string line;
        while (std::getline(f, line)) {
            if (line[0] == '#' || line == "")
                continue;
            std::stringstream ss(line);

            double t_s;
            ss >> t_s;
            data->image_timestamps.emplace_back(t_s * 1e-9);
        }
    }

    void read_imu_data(const std::string &path) {
        std::ifstream fImu;
        fImu.open(path.c_str());

        while (!fImu.eof()) {
            std::string s;
            std::getline(fImu, s);
            if (s[0] == '#' || s == "")
                continue;

            if (!s.empty()) {
                std::string item;
                size_t pos = 0;
                double data_item[7];
                int count = 0;
                while ((pos = s.find(',')) != std::string::npos) {
                    item = s.substr(0, pos);
                    data_item[count++] = std::stod(item);
                    s.erase(0, pos + 1);
                }
                item = s.substr(0, pos);
                data_item[6] = std::stod(item);

                data->imu_timestamps.emplace_back(data_item[0] * 1e-9);
                AccelData::Ptr accel(new AccelData());
                accel->timestamp = data_item[0] * 1e-9;
                accel->data =
                    Eigen::Vector3d(data_item[4], data_item[5], data_item[6]);
                data->accel_data.emplace_back(accel);
                GyroData::Ptr gyro(new GyroData());
                gyro->timestamp = data_item[0] * 1e-9;
                gyro->data =
                    Eigen::Vector3d(data_item[1], data_item[2], data_item[3]);
                data->gyro_data.emplace_back(gyro);
            }
        }
    }

    std::shared_ptr<PanoVideoDataset> data;
};

}  // namespace DatasetIO

#endif